#!/usr/bin/python

'Created by Paul Buckley and Jordan Lange for CSE 415- Intro to AI'
'This program allows a user to train the computer to recognize and'
'guess text that is drawn into a box.'

from Tkinter import *
import sqlite3
import operator
import math
import ast
import os
import heapq
import threading

class Pad:
   def __init__(self):
      self.lastx, self.lasty = 0,0
      self.CANVAS_WIDTH = 900
      self.CANVAS_HEIGHT = 200
      self.PIXELS_PER_GRID = 50
      self.LETTER_DIVISIONS = 8
      self.library = {}
      self.letters = []
      self.currentPadData = [[-1 for col in range(7)] for row in range(8)]
      self.field = [[-1 for col in range(self.CANVAS_WIDTH)] for row in range(self.CANVAS_HEIGHT)]
      self.reducedFieldOffset = []
      self.heatMapList = [] # Keep track of all the heatmaps that have been created
      self.rectangles = [] # Keep track of the segmentations of letters
      self.configureUI()
      ### Load database
      try:
         self.conn = sqlite3.connect(os.path.join(os.path.dirname(__file__), 'learner.db'))
      except NameError:
         self.conn = sqlite3.connect('learner.db')
      self.loadDatabase(self.conn)
      print "Loaded default database: learner"
      
   ############################################################
   ### CALLBACK FUNCTIONS
   ############################################################
   
   #-----------------------------------------------------------
   # Set to the new location on button click
   #-----------------------------------------------------------
   def xy(self, event):
      self.lastx, self.lasty = event.x, event.y
   
   #-----------------------------------------------------------
   # Draws ovals on the canvas where the mouse is and adds that data to the field
   #-----------------------------------------------------------
   def addLine(self, event):
      #if event.x >= 0 and event.y >= 0 and event.x < self.CANVAS_WIDTH and event.y < self.CANVAS_HEIGHT:
         #self.canvas.create_oval((self.lastx, self.lasty, event.x, event.y), width=10)
      for x in range(min(self.lastx, event.x), max(self.lastx, event.x)):
         self.field[event.y][x] = 1
      for y in range(min(self.lasty, event.y), max(self.lasty, event.y)):
         self.field[y][event.x] = 1
      line = self.canvas.create_line(self.lastx, self.lasty, event.x, event.y, width=10, capstyle=ROUND, joinstyle=ROUND)
      self.field[event.y][event.x] = 1
      self.lastx, self.lasty = event.x, event.y
   
   #-----------------------------------------------------------
   # Creates a new database with the text in the entry field
   #-----------------------------------------------------------
   def newDbCallback(self):
      name = self.entry.get()
      if len(name) == 0:
         print "Invalid database name"
      else:
         newConn = createNewDb(name)
         self.loadDatabase(newConn)
         print "Loaded any existing database entries for " + name
      self.entry.delete(0, END)

   #-----------------------------------------------------------
   # Clears the screens and the field data
   #-----------------------------------------------------------
   def clearCallback(self, event=[]):
      self.field = [[-1 for col in range(self.CANVAS_WIDTH)] for row in range(self.CANVAS_HEIGHT)]
      self.clearHeatMaps()
      self.canvas.delete("all")
      self.entry.delete(0, END)
      
   #-----------------------------------------------------------
   # Returns the most likely match to the letter in the current field
   #-----------------------------------------------------------
   def guessCallback(self, event=[]):
      letters = self.getLetterWeights()
   
      # for each letter gotten from getLetterWeights
      # go through dictionary and find best match
      # append best match to a returned value and put it on the self.entry thing by .insert
      
      self.clearHeatMaps()
      guessedLetters = ''
      bestGuesses = []
      guessesToMake = 5
      for letterField in letters:
         bestLetters = {}
         letterWeight = reduce(operator.add, letterField)
         for letter, info in self.library.iteritems():
            t = sum([reduce(lambda x, y: x*y, a) for a in zip(info[1], letterWeight)])
            if (bestLetters.has_key(letter) and t > bestLetters[letter]) or not bestLetters.has_key(letter):
               bestLetters[letter] = t
         ### Create a heap keeping track of the top letters
         if len(bestLetters) == 0:
            print "I haven't learned anything yet!"
            return
         bestHeap = []
         for key in bestLetters.keys():
            heapq.heappush(bestHeap, (bestLetters[key], key))
         ### Get a list of the top letters
         best = heapq.nlargest(guessesToMake, bestHeap)
         ### Add the best letter to the string that goes in the entry field
         guessedLetters += best[0][1]
         ### Keep track of the other guesses
         bestGuesses.append(best[1:guessesToMake])
      ### Draw the heatmaps for each letter
      for i in range(len(guessedLetters)):
         nextLetter = guessedLetters[i]
         location = self.rectangles[i]
         weight = self.library[nextLetter][1]
         letterList = [(nextLetter, weight)]
         ### Build up the other letters and their associated weights
         nextBests = bestGuesses[i]
         for weight, letter in nextBests:
            letterList.append((letter, self.library[letter][1]))
         ### Create the MapContainer object
         mc = MapContainer(self.entry, i, letterList, location, self.LETTER_DIVISIONS)
         #print 'call to mc.createMap(' + nextLetter + ', ' + str(self.library[nextLetter][1]) + ', ' + str(self.LETTER_DIVISIONS) + ', ' + str(location) + ')'
         mc.createMapForGuess(0) ### Guess 0 is the best guess
         self.heatMapList.append(mc)
         mc.show()
      self.entry.delete(0, END)
      self.entry.insert(0, guessedLetters)
   
   #-----------------------------------------------------------
   # Submit the current field data to the dictionary
   #-----------------------------------------------------------
   def submitCallback(self, event=[]):
      # Get the letters from their respective places
      lettersFromField = self.getLetterWeights()
      lettersInText = self.entry.get()
##      for letter in lettersFromField:
##         for row in letter:
##            for item in row:
##               if item == -1:
##                  print " ",
##               else:
##                  print "#",
##            print "\n"
##         print ""
      
      # Get lengths
      lengthText = len(lettersInText)
      lengthFieldData = len(lettersFromField)
      self.clearCallback()
      # If there are a differing number of letters in the field and text box, 
      # tell the user and do nothing else
      if lengthText != lengthFieldData:
         print "According to our calculations, there are ",
         print lengthText,
         print "items in the text box, and ",
         print lengthFieldData,
         print "drawn out.  Try again?"
      # Otherwise we should be able to parse this.
      else:
         # Go through each letter
         for whichLetter in range(0, lengthText):
            fieldLetter = lettersFromField[whichLetter]
            textLetter = lettersInText[whichLetter]
               
            # If the letter in the textbox is not already in the library, add a "blank" entry
            if textLetter not in self.library.keys():
               self.library[textLetter] = [0, [-1]*(self.LETTER_DIVISIONS*self.LETTER_DIVISIONS)]
               
            # Get the data we need from the dict and flatten out the field letter data
            letterDataFromDict = self.library.get(textLetter)
            numTrainings = letterDataFromDict[0]
            letterWeightFromDict = letterDataFromDict[1]
            flattenedFieldLetter = reduce(operator.add, fieldLetter)
               
            # If k is the number of samples you have, weigh the current data by k / k+1
            # and the new data by 1 / k+1
            letterWeightFromDict[:] = [float(x*(numTrainings/(numTrainings+1.0))) for x in letterWeightFromDict] 
            flattenedFieldLetter[:] = [x*(1.0/(numTrainings+1.0)) for x in flattenedFieldLetter]
                  
            # Sum them together to get the new weight
            newWeight = [sum(a) for a in zip(letterWeightFromDict, flattenedFieldLetter)]
               
            # Update that shit (and database!)
            self.library[textLetter] = [ numTrainings+1, newWeight]
            cursor = self.conn.cursor()
            execString = 'insert into data (letter, weights, count) values (\''+str(textLetter)+'\',\''+str(newWeight)+'\','+str(numTrainings+1)+');';
            #print execString
            cursor.execute(execString)
            self.conn.commit()
            cursor.close()

   ############################################################
   ### LETTER SPLICERS
   ############################################################
   
   #-----------------------------------------------------------
   # Returns a LETTER_DIVISIONS x LETTER_DIVISIONS square of weights
   #-----------------------------------------------------------
   def getLetterWeights(self):
      letterWeights = []
      letters = self.getLetters(self.reducedField())
      # Cycle through each letter to get their weights
      for letter in letters:
         wholeLetter = []
         height = len(letter)
         width = len(letter[0])
         
         # We want the weights to be a square so make rectangle into a square
         boxSide = max(height, width)
         pixelsPerDivision = int(math.ceil(float(boxSide) / self.LETTER_DIVISIONS))
               
         # Go through each row and condense the the columns
         temp = []
         for i in range(0, height):
            row = []
            for j in range(0, self.LETTER_DIVISIONS):
               if 1 in letter[i][j*pixelsPerDivision:(j+1)*pixelsPerDivision]:
                     row.append(1)
               else:
                     row.append(-1)
            temp.append(row)
               
         # Now condense the number of rows
         for i in range(0, self.LETTER_DIVISIONS):
            row = []
               
            for j in range(0, self.LETTER_DIVISIONS):
               if 1 in [col[j] for col in temp[i*pixelsPerDivision:(i+1)*pixelsPerDivision]]:
                     row.append(1)
               else:
                     row.append(-1)
            wholeLetter.append(row)
         letterWeights.append(wholeLetter)
         
      return letterWeights
      
   #-----------------------------------------------------------
   # Populates reducedField
   #-----------------------------------------------------------
   def reducedField(self):
      reduced = []
      
      # Zoom in the top and bottom to fit
      top = 0
      bottom = self.CANVAS_HEIGHT-1
      for i in range(0, self.CANVAS_HEIGHT):
         if 1 in self.field[i]:
            top = i
            break
      for i in range(0, self.CANVAS_HEIGHT):
         if 1 in self.field[self.CANVAS_HEIGHT - i - 1]:
            bottom = self.CANVAS_HEIGHT - i
            break
      reduced = self.field[top:bottom]
      
      # Now zoom in the left and right sides
      left = 0
      right = self.CANVAS_WIDTH
      for i in range(0, self.CANVAS_WIDTH):
         if 1 in [row[i] for row in self.field]:
            left = i
            break
      for i in range(0, self.CANVAS_WIDTH):
         if 1 in [row[self.CANVAS_WIDTH - i - 1] for row in self.field]:
            right = self.CANVAS_WIDTH - i
            break
               
      # Return le reduced field
      reduced = [row[left:right] for row in reduced]
      self.reducedFieldOffset = [left, top]
      self.canvas.create_rectangle((left-5, top-5, right+5, bottom+5),width=2)
      return reduced
         
   #-----------------------------------------------------------
   # Returns rectangles that contain each individual letter in the field
   #-----------------------------------------------------------
   def getLetters(self, reducedField):
      letters = []
      fieldWidth = len(reducedField[0])
      fieldHeight = len(reducedField)
      leftEdgePointer = 0
      self.rectangles = [] # Clear any old locations
      
      # Sample from the left, figures should be seperated along the x axis (NOT y)
      while leftEdgePointer < fieldWidth:
         # Initialize the position of the letter
         left, right, top, bottom = 0,0,0,0

         # Found a left edge
         if 1 in [row[leftEdgePointer] for row in reducedField]:
            left = leftEdgePointer
               
            # Find right edge
            for rightEdgePointer in range(left, fieldWidth):
               right = rightEdgePointer
               # Check for right edge by seeing if the next 10 columns are all empty
               if 1 not in sum([row[rightEdgePointer:rightEdgePointer+10] for row in reducedField], []):
                  break
               
            for topEdgePointer in range(0, fieldHeight):
               # Find top edge
               top = topEdgePointer
               if 1 in reducedField[topEdgePointer][left:right+1]:
                  break
                        
            for bottomEdgePointer in range(0, fieldHeight):
               # Find bottom edge
               bottom = fieldHeight - bottomEdgePointer - 1
               if 1 in reducedField[fieldHeight-bottomEdgePointer-1][left:right+1]:
                  break
               
            leftEdgePointer = right+10

            # Get all the data from this run together using the left, right, top, and bottom that were generated
            l = []
            for row in range(top, bottom+1):
               l.append(reducedField[row][left:right+1])

            left = left-5 + self.reducedFieldOffset[0]
            top = top-5 + self.reducedFieldOffset[1]
            right = right+5 + self.reducedFieldOffset[0]
            bottom = bottom+5 + self.reducedFieldOffset[1]
            self.canvas.create_rectangle(tuple([left, top, right, bottom]),width=2) 
            self.rectangles.append(tuple([right-left, bottom-top, left + self.canvas.winfo_rootx(), top + self.canvas.winfo_rooty()])) # width, height, x, y
               
            # Append the letter onto the list of letters
            letters.append(l)
         else:
            leftEdgePointer = leftEdgePointer + 1
      return letters
      
                  
   ############################################################
   ### UI FUNCTIONS
   ############################################################
   def configureUI(self):
      self.root = Tk()
      self.frame = Frame(self.root)
      self.guessButton = Button(self.frame, text="Guess", command=self.guessCallback)
      self.clearButton = Button(self.frame, text="Clear", command=self.clearCallback)
      self.previewMapButton = Button(self.frame, text='Preview\nSubmission', command=self.previewSubmission)
      self.letterFrame = Frame(self.root)
      self.entry = Entry(self.letterFrame)
      self.submitButton = Button(self.letterFrame, text="Submit", command=self.submitCallback)
      self.newDbButton = Button(self.letterFrame, text="Change Db", command=self.newDbCallback)
      self.canvas = Canvas(self.root, width=self.CANVAS_WIDTH, height=self.CANVAS_HEIGHT)
   
      # Configure the root
      self.root.columnconfigure(0, weight=1)
      self.root.rowconfigure(0, weight=1)
      self.root.resizable(False, False)
      self.root.protocol("WM_DELETE_WINDOW", self.exitHandler)

      # Right side buttons
      self.frame.pack(side = RIGHT)
      self.guessButton.pack(side = TOP)
      self.clearButton.pack(side = TOP)
      self.previewMapButton.pack(side = TOP)

      # Bottom bottons
      self.entry.bind("<Return>", self.submitCallback)
      self.entry.config(background="white", highlightthickness=2, highlightcolor="black")
      self.letterFrame.pack(side = BOTTOM)
      self.entry.pack(side = LEFT)
      self.newDbButton.pack(side = RIGHT)
      self.submitButton.pack(side = RIGHT)

      # Main canvas
      self.canvas.grid(column=0, row=0, sticky=(N, W, E, S))
      self.canvas.bind("<Double-Button-1>", self.guessCallback)
      self.canvas.bind("<Button-3>", self.clearCallback)
      self.canvas.bind("<Button-1>", self.xy)
      self.canvas.bind("<Button-2>", self.submitCallback)
      self.canvas.bind("<B1-Motion>", self.addLine)
      self.canvas.config(background="white", borderwidth=5, highlightthickness=4, highlightcolor="green")
      self.canvas.pack(side = LEFT)

   def clearHeatMaps(self):
      for mapRoot in self.heatMapList:
         mapRoot.destroy()
      self.heatMapList = []

   def previewSubmission(self):
      lettersFromField = self.getLetterWeights()
      actualLetters = self.entry.get()
      
      # Copied from submitCallback (Check for input)
      if len(actualLetters) != len(lettersFromField):
         print "According to our calculations, there are ",
         print len(actualLetters),
         print "items in the text box, and ",
         print len(lettersFromField),
         print "drawn out.  Try again?"

      if len(actualLetters) == 0:
         print "Can't draw heatmap -- no input!"
         return
         
      # Clear any old previews
      self.clearHeatMaps()
      
      ## Draw the heatmap for the new input (demonstration purposes, close with right click)
      rIndex = 0
      for letter in lettersFromField:
         location = self.rectangles[rIndex]
         mc = MapContainer(self.entry, rIndex, [(actualLetters[rIndex], reduce(operator.add, letter))], location, self.LETTER_DIVISIONS)
         mc.createMapForGuess(0)
         self.heatMapList.append(mc)
         mc.show()
         rIndex += 1

   def exitHandler(self):
      self.clearHeatMaps()
      self.root.destroy()
      self.conn.close()

   ############################################################
   ### RUN
   ############################################################
   def run(self):
      self.root.mainloop()

   ############################################################
   ### DATABASE
   ############################################################
   def loadDatabase(self, _conn):
      self.library = {}
      self.conn = _conn
      cursor = self.conn.cursor()
      cursor.execute('select MAX(ID), letter, weights, count FROM data\
                                                                     group by letter;')
      for row in cursor:
         key = str(row[1])
         weightList = ast.literal_eval(str(row[2]))
         count = int(row[3])
         self.library[key] = [count, weightList]
      cursor.close()

class MapContainer:
   def __init__(self, entryBar, entryBarOffset, bestLetters, location, divisions):
      self.currentLetter = ''
      self.dims = location
      self.divisions = divisions
      self.bestLetters = bestLetters # bestLetters is an array of the form [(letter, letterGrid), ...]
      ### Keep a handle on the entry bar so we can update it
      self.entryCopy = entryBar
      self.offset = entryBarOffset
      ### Create a color mapping
      self.num_colors = 21
      t = round(255.0/self.num_colors)
      temp = []
      for i in range(self.num_colors):
         temp.append([i, "#%02x%02x%02x" % (t*i, 0, 255-(t*i))])
      self.tk_rgb = dict(temp)
      ### Keep track of each grid
      self.cMapGrid = [] 
      ### Create window and canvas with the properties we need
      self.root = Tk()
      self.root.columnconfigure(0, weight=1)
      self.root.rowconfigure(0,weight=1)
      self.root.resizable(False, False)
      self.root.overrideredirect(True)
      self.root.attributes("-alpha", .7)
      self.canvas = Canvas(self.root)
      self.canvas.pack()
      self.canvas.bind('<Button-3>', lambda x: self.root.withdraw())
      self.canvas.bind('<Button-1>', self.cycleMap)
      self.hide()

   ### Functions to show and hide the window
   def hide(self):
      self.root.withdraw()
      ### Apparantly only works on Windows
      self.root.wm_attributes("-topmost", 0)

   def show(self):
      self.root.update()
      self.root.deiconify()
      ### Apparantly only works on Windows
      self.root.wm_attributes("-topmost", 1)

   def move(self, x, y):
      self.root.geometry('+%s+%s' % (x, y))
                                 
   def destroy(self):
      self.root.destroy()

   ### Creates a color map of  the weights letterGrid
      # gridLocation is an array of the form [x location to place the map, y location to...]
         # These are SCREEN dimensions, not relative window dimensions
      # desiredWidth/desiredHeight are the calculated width/height of the drawn letter
   def createMapForGuess(self, index):
      ### Keep track of the letter we're currently displaying
      self.currentLetter = self.bestLetters[index][0]
      ### REMOVE the rows/columns with all -1
      letterGrid = self.bestLetters[index][1]
      # First convert to a 2d array
      letterGrid = [letterGrid[i:i+self.divisions] for i in range(0, len(letterGrid), self.divisions)]
      # Delete rows
      letterGrid.reverse()
      prev = 0
      while prev != len(letterGrid):
         try:
            letterGrid.remove([-1.0 for j in range(self.divisions)])
            prev = len(letterGrid)
         except ValueError:
            break
      letterGrid.reverse()
      # Delete columns
      remove = True
      while remove:
         for row in letterGrid:
            if row[-1] != -1:
               remove = False
         if remove:
            for row in letterGrid:
               row.pop()
      ### Compute pixels per grid
      xPixelsPerGrid = int(self.dims[0]/len(letterGrid[0]))
      yPixelsPerGrid = int(self.dims[1]/len(letterGrid))
      xStart = self.dims[0]%len(letterGrid[0]) / 2
      yStart = self.dims[1]%len(letterGrid) / 2
      ### Update location of the window
      w = str(self.dims[0])
      h = str(self.dims[1])
      x = str(self.dims[2])
      y = str(self.dims[3])
      self.root.geometry(w+"x"+h+"+"+x+"+"+y)
      ### Draw grid
      y = yStart
      for row in letterGrid:
         x = xStart
         for col in row:
            self.cMapGrid.append\
            (\
               self.canvas.create_rectangle\
               (\
                  x, y, x+xPixelsPerGrid, y+yPixelsPerGrid, \
                  fill=self.tk_rgb[int((col+1.0)*10%10) + 10*int(col+1.0)]\
               )\
            )
            x += xPixelsPerGrid
         y += yPixelsPerGrid
   
   #-----------------------------------------------------------
   # Callback to handle displaying the next best guess's heatmap!
   #-----------------------------------------------------------
   def cycleMap(self, event):
      if len(self.bestLetters) == 1:
         return # Can't cycle on only 1 letter
      for i in range(len(self.bestLetters)):
         if self.bestLetters[i][0] == self.currentLetter:
            nextIndex = (i+1) % (len(self.bestLetters)-1)
            ### Clear the canvas and draw the next letter
            self.clearCanvas()
            self.createMapForGuess(nextIndex)
            ### Update the entry bar
            self.entryCopy.delete(self.offset, self.offset+1)
            self.entryCopy.insert(self.offset, self.bestLetters[nextIndex][0])
            return

   ### Clears the color map on the grid
   def clearCanvas(self):
      for subGrid in self.cMapGrid:
         self.canvas.delete(subGrid)

def createNewDb(name):
   try:
      conn = sqlite3.connect(os.path.join(os.path.dirname(__file__), name+'.db'))
   except NameError:
      conn = sqlite3.connect(name + '.db')
   cursor = conn.cursor()
   execString = 'CREATE TABLE data (ID INTEGER PRIMARY KEY AUTOINCREMENT, \
                     letter TEXT, \
                     weights TEXT, \
                     count int);'
   try:
      cursor.execute(execString)
      conn.commit()
      print 'Created new database: ' + name
   except sqlite3.OperationalError:
      print 'Database ' + name + ' already exists. '
   cursor.close()
   return conn

if __name__ == "__main__":
   pad = Pad()
   pad.run()
